import os 

class ABE:

    #Note：会生成对应文件，可以输出二进制内容
    @Output(pub_key, master_key)
    def setup(self):
        os.system("./cpabe-setup")

        file = open('pub_key',"rb")
        pub_key = file.read()
        file.close()

        file = open('master_key',"rb")
        master_key = file.read()
        file.close()

        return pub_key,master_key
    
    #Problem：pk和mk只能输入文件名而非内容；私钥会生成对应文件，可以输出二进制内容
    @Input(user_attr)
    @Output(user_priv_key)
    def keygen(self,user_attr):
        os.system("./cpabe-keygen -o user_priv_key pub_key master_key " + user_attr)

        file = open('user_priv_key',"rb")
        user_priv_key = file.read()
        file.close()
        
        return user_priv_key

    #Problem：pk和mk只能输入文件名而非内容；密文会生成对应文件，可以输出二进制内容
    @Input(msg,access_policy)
    @Output(ciphertext)
    def encrypt(self,msg,access_policy):
        file = open("user_security_file","w")
        file.write(msg)
        file.close()

        os.system("./cpabe-enc pub_key user_security_file " + access_policy)

        file = open("user_security_file.cpabe","rb")
        ciphertext = file.read()
        file.close()

        return ciphertext

    #Problem：pk和用户私钥和加密后的文件只能输入文件名而非内容；明文会生成对应文件
    @Output(msg)
    def decrypt(self):
        os.system("./cpabe-dec pub_key user_priv_key user_security_file.cpabe")

        file = open("user_security_file")
        msg = file.read()
        file.close()  

        return msg


if __name__ == '__main__':

    abeObj = ABE()

    pub_key,master_key = abeObj.setup()
    print("公共秘钥 =>   ",pub_key)
    print("\n\n\n\n\n")
    print("主密钥 =>   ",master_key)
    print("\n\n\n\n\n")
    
    user_priv_key = abeObj.keygen("sysadmin it_department")
    print("用户私钥 =>   ",user_priv_key)
    print("\n\n\n\n\n")

    ciphertext = abeObj.encrypt("123456abcdef","sysadmin")
    print("加密后的密文 =>  ",ciphertext)
    print("\n\n\n\n\n")

    msg = abeObj.decrypt()
    print("解密后的明文 =>  ",msg)



